import numpy as np
import itertools
from Xponge.helper import AbstractMolecule, Molecule, Xdict
from Xponge.build import build_bonded_force
from .molecule import Molecule as tMolecule
from ..potential import ForceFieldBase

Molecule._torchsponge_todo = Xdict()

@classmethod
def set_torchsponge_todo(cls, keyname):

    def wrapper(func):
        cls._torchsponge_todo[keyname] = func
        return func

    return wrapper

@classmethod
def del_torchsponge_todo(cls, keyname):

    cls._torchsponge_todo.pop(keyname)

Molecule.set_torchsponge_todo = set_torchsponge_todo
Molecule.del_torchsponge_todo = del_torchsponge_todo

@Molecule.set_torchsponge_todo("coordinate")
def _do(self, sys_kwarg, ene_kwarg, use_pbc):
    if self.box_length is not None:
        box_length = self.box_length/10
    else:
        crd = self.get_atom_coordinates()/10
        box_length = np.array(np.max(crd, axis=0) - np.min(crd, axis=0) + 0.6, dtype=np.float32).tolist()
    if "coordinate" not in sys_kwarg:
        crd = self.get_atom_coordinates()/10
        sys_kwarg["coordinate"] = [crd.tolist()]
        sys_kwarg["atom_name"] = [[atom.name for atom in self.atoms]]
        sys_kwarg["atom_type"] = [[atom.type.name for atom in self.atoms]]
        sys_kwarg["pbc_box"] = [box_length]
        sys_kwarg["residue_name"] = [[residue.name for residue in self.residues]]
        index = [0] + [len(self.residues[i].atoms) for i in range(0,len(self.residues))][:-1]
        residue_index = list(itertools.accumulate(index))
        res_id = [i for i in range(0,len(self.residues)) for j in range(0,len(self.residues[i].atoms))]
        sys_kwarg["residue_index"] = [res_id]
    else:
        crd = self.get_atom_coordinates()/10
        sys_kwarg["coordinate"].append(crd.tolist())
        sys_kwarg["atom_name"].append([atom.name for atom in self.atoms])
        sys_kwarg["atom_type"] = [[atom.type.name for atom in self.atoms]]
        sys_kwarg["pbc_box"].append(box_length)
        sys_kwarg["residue_name"] = [[residue.name for residue in self.residues]]
        index = [0] + [len(self.residues[i].atoms) for i in range(0,len(self.residues))][:-1]
        residue_index = list(itertools.accumulate(index))
        res_id = [i for i in range(0,len(self.residues)) for j in range(0,len(self.residues[i].atoms))]
        sys_kwarg["residue_index"] = [res_id]


@Molecule.set_torchsponge_todo("charge")
def _do(self, sys_kwarg, ene_kwarg, use_pbc):

    from ..potential import CoulombEnergy
    if "atom_charge" not in sys_kwarg:
        sys_kwarg["atom_charge"] = []
    sys_kwarg["atom_charge"].append([atom.charge for atom in self.atoms])
    if "charge" not in ene_kwarg:
        ene_kwarg["charge"] = Xdict()
        if use_pbc:
            ene_kwarg["charge"]["function"] = lambda system, ene_kwarg: CoulombEnergy(
                atom_charge=system.atom_charge, pbc_box=system.pbc_box, use_pme=True,
                nfft=system.pbc_box.detach().cpu().numpy().astype(int)//4*4,
                exclude_index=np.array(sys_kwarg["exclude"]))
        else:
            ene_kwarg["charge"]["function"] = lambda system, ene_kwarg: CoulombEnergy(
                atom_charge=system.atom_charge, exclude_index=np.array(sys_kwarg["exclude"]))

  
@Molecule.set_torchsponge_todo("mass")
def _do(self, sys_kwarg, ene_kwarg, use_pbc):
    """
    :return:
    """
    if "atom_mass" not in sys_kwarg:
        sys_kwarg["atom_mass"] = []
    sys_kwarg["atom_mass"].append([atom.mass for atom in self.atoms])


@Molecule.set_torchsponge_todo("bond")
def _do(self, sys_kwarg, ene_kwarg, use_pbc):
 
    from ..potential import BondEnergy
    if "bonds" not in sys_kwarg:
        sys_kwarg["bonds"] = []
    if "bond" not in ene_kwarg:
        ene_kwarg["bond"] = Xdict()
        ene_kwarg["bond"]["function"] = lambda system, ene_kwarg: BondEnergy(
            index=system.bonds, use_pbc=use_pbc,
            force_constant=ene_kwarg["bond"]["force_constant"],
            bond_length=ene_kwarg["bond"]["bond_length"],
            )
        ene_kwarg["bond"]["force_constant"] = []
        ene_kwarg["bond"]["bond_length"] = []
    bonds = []
    force_constants = []
    bond_lengths = []
    for bond in self.bonded_forces.get("bond", []):
        if bond.k == 0:
            continue
        bonds.append([self.atom_index[bond.atoms[0]], self.atom_index[bond.atoms[1]]])
        force_constants.append(bond.k * 2 * 4.18 / 0.01)
        bond_lengths.append(bond.b * 0.1)
    sys_kwarg["bonds"].append(bonds)
    ene_kwarg["bond"]["force_constant"].append(force_constants)
    ene_kwarg["bond"]["bond_length"].append(bond_lengths)
            

@Molecule.set_torchsponge_todo("angle")
def _do(self, sys_kwarg, ene_kwarg, use_pbc):

    from ..potential import AngleEnergy
    if "angle" not in ene_kwarg:
        ene_kwarg["angle"] = Xdict()
        ene_kwarg["angle"]["function"] = lambda system, ene_kwarg: AngleEnergy(
            index=ene_kwarg["angle"]["index"], use_pbc=use_pbc,
            force_constant=ene_kwarg["angle"]["force_constant"],
            bond_angle=ene_kwarg["angle"]["bond_angle"],
            )
        ene_kwarg["angle"]["index"] = []
        ene_kwarg["angle"]["force_constant"] = []
        ene_kwarg["angle"]["bond_angle"] = []
    bonds = []
    force_constants = []
    bond_angles = []
    for bond in self.bonded_forces.get("angle", []):
        if bond.k == 0:
            continue
        bonds.append([self.atom_index[atom] for atom in bond.atoms])
        force_constants.append(bond.k * 2 * 4.18)
        bond_angles.append(bond.b)
    ene_kwarg["angle"]["index"].append(bonds)
    ene_kwarg["angle"]["force_constant"].append(force_constants)
    ene_kwarg["angle"]["bond_angle"].append(bond_angles)


@Molecule.set_torchsponge_todo("dihedral")
def _do(self, sys_kwarg, ene_kwarg, use_pbc):

    from ..potential import DihedralEnergy
    if "dihedral" not in ene_kwarg:
        ene_kwarg["dihedral"] = Xdict()
        ene_kwarg["dihedral"]["function"] = lambda system, ene_kwarg: DihedralEnergy(
            index=ene_kwarg["dihedral"]["index"], use_pbc=use_pbc,
            force_constant=ene_kwarg["dihedral"]["force_constant"],
            periodicity=ene_kwarg["dihedral"]["periodicity"],
            phase=ene_kwarg["dihedral"]["phase"],
            )
        ene_kwarg["dihedral"]["index"] = []
        ene_kwarg["dihedral"]["force_constant"] = []
        ene_kwarg["dihedral"]["periodicity"] = []
        ene_kwarg["dihedral"]["phase"] = []
    dihedrals = []
    force_constants = []
    periodicitys = []
    phases = []
    for dihedral in self.bonded_forces.get("dihedral", []):
        for i in range(dihedral.multiple_numbers):
            if dihedral.ks[i] != 0:
                dihedrals.append([self.atom_index[atom] for atom in dihedral.atoms])
                force_constants.append(dihedral.ks[i] * 2 * 4.18)
                periodicitys.append(dihedral.periodicitys[i])
                phases.append(dihedral.phi0s[i])
    for dihedral in self.bonded_forces.get("improper", []):
        if dihedral.k != 0:
            dihedrals.append([self.atom_index[atom] for atom in dihedral.atoms])
            force_constants.append(dihedral.k * 2 * 4.18)
            periodicitys.append(dihedral.periodicity)
            phases.append(dihedral.phi0)
    ene_kwarg["dihedral"]["index"].append(dihedrals)
    ene_kwarg["dihedral"]["force_constant"].append(force_constants)
    ene_kwarg["dihedral"]["periodicity"].append(periodicitys)
    ene_kwarg["dihedral"]["phase"].append(phases)



@Molecule.set_torchsponge_todo("LJ")
def _do(self, sys_kwarg, ene_kwarg, use_pbc):

    from Xponge.forcefield.base.lj_base import LJType
    from ..potential import LennardJonesEnergy
    if "lj" not in ene_kwarg:
        ene_kwarg["lj"] = Xdict()
        ene_kwarg["lj"]["function"] = lambda system, ene_kwarg: LennardJonesEnergy(
            epsilon=ene_kwarg["lj"]["epsilon"], use_pbc=use_pbc,
            sigma=ene_kwarg["lj"]["sigma"],)
        ene_kwarg["lj"]["epsilon"] = []
        ene_kwarg["lj"]["sigma"] = []
    ljtypes = [LJType.get_type(atom.LJtype + "-" + atom.LJtype) for atom in self.atoms]
    ene_kwarg["lj"]["epsilon"].append([ljtype.epsilon * 4.18 for ljtype in ljtypes])
    ene_kwarg["lj"]["sigma"].append([ljtype.rmin / (4 ** (1 / 12) / 2) * 0.1 for ljtype in ljtypes])


@Molecule.set_torchsponge_todo("nb14")
def _do(self, sys_kwarg, ene_kwarg, use_pbc):

    from ..potential import NonbondPairwiseEnergy
    from Xponge.forcefield.base.nb14_extra_base import get_nb14_extra_lj
    if "nb14" not in ene_kwarg:
        ene_kwarg["nb14"] = Xdict()
        ene_kwarg["nb14"]["function"] = lambda sys_kwarg, ene_kwarg: NonbondPairwiseEnergy(
            index=ene_kwarg["nb14"]["index"], use_pbc=use_pbc,
            qiqj=ene_kwarg["nb14"]["qiqj"],
            epsilon_ij=ene_kwarg["nb14"]["epsilon_ij"],
            sigma_ij=ene_kwarg["nb14"]["sigma_ij"],
            r_scale=ene_kwarg["nb14"]["r_scale"],
            r6_scale=ene_kwarg["nb14"]["r6_scale"],
            r12_scale=ene_kwarg["nb14"]["r12_scale"],
            )
        ene_kwarg["nb14"]["index"] = []
        ene_kwarg["nb14"]["r_scale"] = []
        ene_kwarg["nb14"]["r6_scale"] = []
        ene_kwarg["nb14"]["r12_scale"] = []
        ene_kwarg["nb14"]["qiqj"] = []
        ene_kwarg["nb14"]["sigma_ij"] = []
        ene_kwarg["nb14"]["epsilon_ij"] = []
    nb14s = []
    r_scales = []
    r6_scales = []
    r12_scales = []
    qiqj = []
    sigma_ij = []
    epsilon_ij = []
    for bond in self.bonded_forces.get("nb14", []):
        if bond.kee != 1 and bond.kLJ != 1:
            atom1, atom2 = bond.atoms
            a, b = get_nb14_extra_lj(atom1, atom2)
            nb14s.append([self.atom_index[atom] for atom in bond.atoms])
            qiqj.append(atom1.charge * atom2.charge)
            if a != 0:
                sigma_ij.append((a/b)**(1/6) * 0.1)
                epsilon_ij.append(b*b/a/4 * 4.18)
            else:
                sigma_ij.append(0)
                epsilon_ij.append(0)
            r_scales.append(bond.kee)
            r6_scales.append(bond.kLJ)
            r12_scales.append(bond.kLJ)
    ene_kwarg["nb14"]["index"].append(nb14s)
    ene_kwarg["nb14"]["r_scale"].append(r_scales)
    ene_kwarg["nb14"]["r6_scale"].append(r6_scales)
    ene_kwarg["nb14"]["r12_scale"].append(r12_scales)
    ene_kwarg["nb14"]["qiqj"].append(qiqj)
    ene_kwarg["nb14"]["sigma_ij"].append(sigma_ij)
    ene_kwarg["nb14"]["epsilon_ij"].append(epsilon_ij)

def _get_single_system_energy(scls, sys_kwarg, ene_kwarg, use_pbc):

    for todo in getattr(scls, "_torchsponge_todo").values():
        todo(scls, sys_kwarg, ene_kwarg, use_pbc)
    getattr(scls, "_mindsponge_todo")['exclude'](scls, sys_kwarg, ene_kwarg, use_pbc)

def get_torchsponge_energy(cls, cutoff=None, **kwargs):
    from collections.abc import Iterable
    use_pbc = True
    if cutoff is None:
        print("[Warning] Cutoff is None, no pbc will be used!")
        use_pbc = False
    if not isinstance(cls, Iterable):
        cls = [cls]
    sys_kwarg = Xdict()
    ene_kwarg = Xdict()
    for scls in cls:
        if isinstance(scls, AbstractMolecule):
            mol = Molecule.cast(scls)
            build_bonded_force(mol)
            _get_single_system_energy(mol, sys_kwarg, ene_kwarg, use_pbc)
        else:
            raise TypeError(f"The type should be a Molecule, Residue, ResidueType, but we get {str(type(scls))}")
    toremove = []
    for key, value in sys_kwarg.items():
        if not value[0]:
            toremove.append(key)
        if key == "pbc_box" and cutoff == None:
            toremove.append(key)
    for key in toremove:
        sys_kwarg.pop(key)

    system = tMolecule(**sys_kwarg)
    system.multi_system = len(cls)
    energies = []
    sys_kwarg["exclude"] = ene_kwarg.pop("exclude")
    for todo in ene_kwarg.values():
        try:
            energies.append(todo["function"](system, ene_kwarg))
        except (TypeError, ValueError) as e:
            if 'NoneType' not in e.args[0] and 'zero dimension' not in e.args[0] and 'but got: 0' not in e.args[0]:
                raise e

    try:
        energy = ForceFieldBase(energy=energies, cutoff=cutoff, exclude_index=sys_kwarg["exclude"])
    except ValueError as e:
        if 'zero dimension' not in e.args[0]:
            raise e
        energy = ForceFieldBase(energy=energies)

    return system, energy

